/*
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov;

import java.util.HashMap;
import java.util.Map;

public class DecayingLearningRate implements LearningRate {
	
	protected Map<State, int[]> visitCount;
	
	protected int numOptions, temp[];
	
	protected double start, exponent;
	
	/**
	 * In this implementation learning rate for a state or state-optionpair is decayed by
	 * 1 / t^w, where t is the number of visits to the state(-option) pair.
	 * 
	 * This constructor sets w = 0.5.
	 * 
	 * @param numOptions size of the option space
	 */
	public DecayingLearningRate(int numOptions) {
		this(1, numOptions, 0.5);
	}
	
	/**
	 * In this implementation learning rate for a state or state-optionpair is decayed by
	 * start / t^w, where t is the number of visits to the state(-option) pair.
	 * 
	 * @param numOptions size of the option space
	 * @param exponent the w in above formula
	 */
	public DecayingLearningRate(double start, int numOptions, double exponent) {
		visitCount = new HashMap<State, int[]>();
		this.numOptions = numOptions;
		this.exponent = exponent;
		this.start = start;
	}
	
	public double getExponent() {
		return exponent;
	}

	public void setExponent(double exponent) {
		this.exponent = exponent;
	}
	
	public <T> void visit(State<T> s) {
		if ( numOptions > 1 )
			throw new UnsupportedOperationException("Number of options is greater than 1 (" + numOptions + "). Use visit(State,Option) instead.");
		
		temp = visitCount.get(s);
		if ( temp == null ) {
			temp = new int[1];
			visitCount.put(s, temp);
		}
		temp[0]++;
	}
	
	public <T> void visit(State<T> s, Option o) {
		temp = visitCount.get(s);
		if ( temp == null ) {
			temp = new int[numOptions];
			visitCount.put(s, temp);
		}
		temp[o.getID()]++;
	}
	
	public <T> double get(State<T> s) {
		return start / Math.pow(getCount(s)+1, exponent);
	}
	
	public <T> double get(State<T> s, Option o) {
		return start / Math.pow(getCount(s,o)+1, exponent);
	}
	
	public void init() {
		visitCount.clear();
	}
	
	protected <T> int getCount(State<T> s) {
		temp = visitCount.get(s);
		if ( temp == null )
			return 0;
		int cnt = 0;
		for ( int i = 0; i < numOptions; i++ )
			cnt += temp[i];
		return cnt;
	}
	
	protected <T> int getCount(State<T> s, Option o) {
		temp = visitCount.get(s);
		if ( temp == null )
			return 0;
		return temp[o.getID()];
	}

}
